#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@File    : btree.py
@Time    : 2022/11/25 18:49:42
@Author  : 郭瑞强
@Contact : sunraing@126.com
@Version : 0.1
@License : BSD 3-Clause License
@Desc    : 利用链表表实现二叉树
"""

from pythonds.basic import Queue


class BTree:
    def __init__(self, root_obj=None, left=None, right=None):
        self.__data = root_obj
        self.__left_child: BTree = left
        self.__right_child: BTree = right

    def insert_left(self, new_node):
        if self.__left_child is None:
            self.__left_child = BTree(new_node)
        else:
            t = BTree(new_node)
            t.__left_child = self.__left_child
            self.__left_child = t

    def insert_right(self, new_node):
        if self.__right_child is None:
            self.__right_child = BTree(new_node)
        else:
            t = BTree(new_node)
            t.__right_child = self.__right_child
            self.__right_child = t

    def get_left_child(self):
        return self.__left_child

    def get_right_child(self):
        return self.__right_child

    def get_root_val(self):
        return self.__data

    def set_root_val(self, obj):
        self.__data = obj

    def preorder(self):
        print(self.__data)
        if self.__left_child:
            self.__left_child.preorder()
        if self.__right_child:
            self.__right_child.preorder()

    def inorder(self):
        if self.__left_child:
            self.__left_child.inorder()
        print(self.__data)
        if self.__right_child:
            self.__right_child.inorder()

    def postorder(self):
        if self.__left_child:
            self.__left_child.postorder()
        if self.__right_child:
            self.__right_child.postorder()
        print(self.__data)

    def layer_order(self):
        node_queue = Queue()
        node_queue.enqueue(self)
        while not node_queue.empty():
            node: BTree = node_queue.dequeque()
            print(node.get_root_val())
            if node.get_left_child():
                node_queue.enqueue(node.get_left_child())
            if node.get_right_child():
                node_queue.enqueue(node.get_right_child())

    def layer_count(self) -> int:
        layer_list = []  # 按层存储节点
        node_queue = [self]  # 存储当前层的所有节点

        while node_queue:
            next_layer_nodes = []  # 下一层的所有节点
            used_node_list = []  # 已经使用过的节点的值
            while node_queue:
                node: BTree = node_queue.pop(0)
                if node.get_left_child():
                    next_layer_nodes.append(node.get_left_child())
                if node.get_right_child():
                    next_layer_nodes.append(node.get_right_child())
                used_node_list.append(node.get_root_val())

            node_queue = next_layer_nodes
            layer_list.append(used_node_list)

        # print(layer_list)
        return len(layer_list)


from pythonds.node import TSNode


class BSTree:
    """二分查找树:dict的一种实现"""

    def __init__(self) -> None:
        self.root: TSNode = None
        self.__size = 0

    def size(self):
        return self.__size

    def __len__(self):
        return self.__size

    def __iter__(self):
        return self.root.__iter__()

    # put --------------------------------------------------------------------
    def put(self, key, val):
        if self.root:
            node = self._get(key, self.root)
            if node:
                node.payload = val
            else:
                self._put(key, val, self.root)
                self.__size += 1
        else:
            self.root = TSNode(key, val)
            self.__size += 1

    def _put(self, key, val, cur_node: TSNode):
        if key < cur_node.key:
            if cur_node.has_left_child():
                self._put(key, val, cur_node.left_child)
            else:
                cur_node.left_child = TSNode(key, val, None, None, parent=cur_node)
        else:
            if cur_node.has_right_child():
                self._put(key, val, cur_node.right_child)
            else:
                cur_node.right_child = TSNode(key, val, None, None, parent=cur_node)

    def __setitem__(self, key, val):
        self.put(key, val)

    # get --------------------------------------------------------------------
    def get(self, key):
        if self.root:
            res: TSNode = self._get(key, self.root)
            if res:
                return res.payload
            else:
                return None

        return None

    def _get(self, key, cur_node: TSNode):
        if not cur_node:
            return None

        if cur_node.key == key:
            return cur_node

        if cur_node.key > key:
            return self._get(key, cur_node.left_child)

        return self._get(key, cur_node.right_child)

    def __getitem__(self, key):
        return self.get(key)

    # in ---------------------------------------------------------------------
    def __contains__(self, key):
        if self._get(key, self.root):
            return True

        return False

    # delete -----------------------------------------------------------------
    def __delitem__(self, key):
        self.delete(key)

    def delete(self, key):
        if self.__size > 1:
            node_to_del = self._get(key, self.root)
            if node_to_del:
                self.remove(node_to_del)
                self.__size -= 1
            else:
                raise KeyError("错误, 键值不存在。")

        elif self.__size == 1 and self.root.key == key:
            self.root = None
            self.__size -= 1
        else:
            raise KeyError("错误, 键值不存在。")

    def remove(self, cur_node: TSNode):

        if cur_node.is_leaf():  # 叶子节点
            # 如果该节点是根节点，则只有一个节点，此情况在delete中已经实现，不用再考虑
            # 即需要考虑的节点都是非根节点，肯定有父节点
            self.leaf_del(cur_node)

        elif cur_node.has_both_child():  # 有双子点节
            # 不能简单的将某个子点上移替换被删节点
            # 被删节点右子树中最小的那个节点，即为“后继”节点：
            # 后继节点最多只有1个：本身是叶子节点，或只有唯一的右子树
            # 将这个后继节点摘出来（也就是删除了），替换掉被 删除节点
            self.twochild_del(cur_node)

        else:  # 只有一个子节点情况, 子节点上移
            self.onechild_del(cur_node)

    def leaf_del(self, cur_node: TSNode):
        # parent_info 当前节点是其父节点的哪个子节点
        parent_info = "left_child" if cur_node.is_left_child() else "right_child"
        setattr(cur_node.parent, parent_info, None)

    def onechild_del(self, cur_node: TSNode):
        # child_info 当前节点有哪个子节点
        child_info = "left_child" if cur_node.has_left_child() else "right_child"
        child: TSNode = getattr(cur_node, child_info)
        if cur_node.is_root():
            cur_node.replace_node_data(
                child.key, child.payload, child.left_child, child.right_child
            )
        else:
            # parent_info 当前节点是其父节点的哪个子节点
            parent_info = "left_child" if cur_node.is_left_child() else "right_child"
            setattr(cur_node.parent, parent_info, child)
            child.parent = cur_node.parent

    def twochild_del(self, cur_node: TSNode):
        # 寻找后继
        succ: TSNode = self.find_successor(cur_node)

        # 后继节点属性替换当前节点
        cur_node.key = succ.key
        cur_node.payload = succ.payload

        # 删除后继节点
        # 找到的后继节点只可能是叶子节点或只有一个子节点
        if succ.is_leaf():
            self.leaf_del(succ)
        else:
            self.onechild_del(succ)

    def find_successor(self, node: TSNode):
        # 拥有2个子节点的删除，需要找到当前节点的后继
        # 从右子树找最小值
        succ = node.right_child
        while succ.has_left_child():
            succ = succ.left_child

        return succ


"""AVL平衡二叉树
改进BST的性能,避免退化情形
作为BST,新key必定以**叶节点**形式插入到AVL树中

叶节点的平衡因子是0,其本身无需重新平衡,但会影响其父节点的平衡因子:
- 作为左子节点插入,则父节点平衡因子会加1
- 作为右子节点插入,则父节点平衡因子会减1

这种影响可能随着其父节点到根节点的路径一直传递上去，直到：
- 传递到根节点为止；
- 或者某个父节点平衡因子被调整到0,不再影响上层节点的平衡因子为止
    无论从-1或者1调整到0,都不会改变子树的高度
"""


class AVLT(BSTree):
    """AVL平衡二叉树"""

    def _put(self, key, val, cur_node: TSNode):
        if key < cur_node.key:
            if cur_node.has_left_child():
                self._put(key, val, cur_node.left_child)
            else:
                cur_node.left_child = TSNode(key, val, None, None, cur_node)
                self.update_balance(cur_node.left_child)
        else:
            if cur_node.has_right_child():
                self._put(key, val, cur_node.right_child)
            else:
                cur_node.right_child = TSNode(key, val, None, None, cur_node)
                self.update_balance(cur_node.right_child)

    def update_balance(self, node: TSNode, op="put"):
        if node.balance_factor > 1 or node.balance_factor < -1:
            self.rebalance(node)  # 旋转节点

            if op == "put":
                # put后的重新平衡，子树高度不变，新节点平衡因子为0,终止向上传播,
                return
            else:
                # del后的重新平衡，子树高度减1，新节点平衡因子为0,继续向上传播
                # 重新平衡后，node节点下沉，其子节点上浮
                node = node.parent

        if node.parent is not None:
            if (op == "put" and node.is_left_child()) or (
                op == "del" and node.is_right_child()
            ):
                node.parent.balance_factor += 1
            else:
                node.parent.balance_factor -= 1

            if op == "put":
                # 插入后，子树高度加1,向上传递，直至根节点或平衡节点
                if node.parent.balance_factor != 0:
                    self.update_balance(node.parent, op)
            else:
                # 删除后，没有了子点，子树高度降1，需要向上传递,直至根节点或深度不变节点
                if node.parent.balance_factor not in {1, -1}:
                    self.update_balance(node.parent, op)

    def rebalance(self, node: TSNode):
        """重新平衡
        主要手段：将不平衡的子树进行**旋转**，rotation
        分“左重”或者“右重进行不同方向的旋转，
        同时更新相关父节点的引用，
        更新旋转后被影响节点的平衡因子
        """
        if node.balance_factor < 0:
            if node.right_child.balance_factor > 0:
                self.rotate_right(node.right_child)
            self.rotate_left(node)
        elif node.balance_factor > 0:
            if node.left_child.balance_factor < 0:
                self.rotate_left(node.left_child)
            self.rotate_right(node)

    def rotate_left(self, rot_root: TSNode):
        # 新的根节点
        new_root = rot_root.right_child

        # 新根节点的左子节点调整到rot的右子节点
        rot_root.right_child = new_root.left_child
        if new_root.has_left_child():
            new_root.left_child.parent = rot_root

        # 新根节点的父节点调整为旧根的父节点
        new_root.parent = rot_root.parent
        if rot_root.is_root():
            self.root = new_root
        else:
            if rot_root.is_left_child():
                rot_root.parent.left_child = new_root
            else:
                rot_root.parent.right_child = new_root

        # 旧根调整为新根的左子节点
        new_root.left_child = rot_root
        rot_root.parent = new_root

        rot_root.balance_factor = (
            rot_root.balance_factor + 1 - min(new_root.balance_factor, 0)
        )
        new_root.balance_factor = (
            new_root.balance_factor + 1 + max(rot_root.balance_factor, 0)
        )

    def rotate_right(self, rot_root: TSNode):
        # 新的根节点
        new_root = rot_root.left_child

        # 新根节点的右子节点调整为旧根的左子节点
        rot_root.left_child = new_root.right_child
        if new_root.has_right_child():
            new_root.right_child.parent = rot_root

        # 新根节点的父节点调整为旧根的父子点
        new_root.parent = rot_root.parent
        if rot_root.is_root():
            self.root = new_root
        else:
            if rot_root.is_left_child():
                rot_root.parent.left_child = new_root
            else:
                rot_root.parent.right_child = new_root

        # 旧根调整为新根的右子节点
        new_root.right_child = rot_root
        rot_root.parent = new_root

        rot_root.balance_factor = (
            rot_root.balance_factor - 1 - max(new_root.balance_factor, 0)
        )
        new_root.balance_factor = (
            new_root.balance_factor - 1 + min(rot_root.balance_factor, 0)
        )

    def balance_del(fn):
        def decorator(self, *args):
            cur_node: TSNode = args[0]
            parent = cur_node.parent
            if parent:
                if cur_node.is_left_child():
                    parent.balance_factor -= 1
                else:
                    parent.balance_factor += 1

            # self.fn(cur_node)
            fn(self, *args)

            if parent and parent.balance_factor not in {1, -1}:
                self.update_balance(parent, "del")

        return decorator

    @balance_del
    def leaf_del(self, cur_node: TSNode):
        # 删除节点前先调整父节点的平衡因子
        # 与插入节点相反，删除左子节点，平衡因子减1；删除右子节点，平衡因子加1
        # parent = cur_node.parent
        # if cur_node.is_left_child():
        #     parent.balance_factor -= 1
        # else:
        #     parent.balance_factor += 1

        # 删除节点
        super().leaf_del(cur_node)

        # parent.balance_factor
        # -1, 1 : 有两个子节点，删除了一个，子树深度不变，不向上传播
        # 0：只有一个节点，删除后子树深度减1，向上传递，直到根节点或平衡因子为1或-1
        # 2, -2：需要重新平衡，重新平衡后子树深度减速1,向上传递，直到根节点或平衡因子为1或-1
        # if parent.balance_factor not in {1, -1}:
        #     self.update_balance(parent, "del")

    @balance_del
    def onechild_del(self, cur_node: TSNode):
        # 该节点只有一个叶子节点
        # parent = cur_node.parent
        # if parent:
        #     if cur_node.is_left_child():
        #         parent.balance_factor -= 1
        #     else:
        #         parent.balance_factor += 1

        # 删除节点
        super().onechild_del(cur_node)

        # if parent and parent.balance_factor not in {1, -1}:
        #     self.update_balance(parent, "del")

    # 内部调用的移除节点函数已经装饰过，此处不用再装饰
    # def twochild_del(self, cur_node: TSNode):
    #     # 找到后继节点，替换key及存储信息
    #     succ: TSNode = self.find_successor(cur_node)
    #     cur_node.key = succ.key
    #     cur_node.payload = succ.payload

    #     # 删除后继节点
    #     # 找到的后继节点只可能是叶子节点或只有一个子节点
    #     if succ.is_leaf():
    #         self.leaf_del(succ)
    #     else:
    #         self.onechild_del(succ)
    # 和父类内容一致
    # 涉及到节点删除会调用AVLT类的方法，会自动加装饰，没必要再实现一次

    def find_successor(self, node: TSNode):
        # 拥有2个子节点的删除，需要找到当前节点的后继
        succ: TSNode = None

        # 当前节点的平衡因子可能是1、 0、-1
        # 当左子树更深时（即平衡因子为1时），从左子树找最大值
        # 否则从右子树找最小值
        if node.balance_factor == 1:
            succ = node.left_child
            while succ.has_right_child():
                succ = succ.right_child
        else:
            succ = node.right_child
            while succ.has_left_child():
                succ = succ.left_child

        return succ
